/*
Copyright 2017 Google Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS-IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

#include "seurat/tiler/selection_tiler.h"

#include <map>
#include <numeric>
#include <vector>

#include "absl/types/span.h"
#include "seurat/base/parallel.h"
#include "seurat/base/reporting.h"
#include "seurat/tiler/selection/selection_problem.h"
#include "seurat/tiler/subdivision_util.h"

namespace seurat {
namespace tiler {

using selection::Token;

namespace {

// Each cell of the Subdivision has some number of possible tilings which may
// be used.
//
// This struct stores the relationship between the SelectionSolver's 'item'
// representing a set of 'tiles' generated by the CandidateTileGenerator.
struct Candidate {
  const CandidateTileGenerator::CandidateTiles* tiles;
  int item;
  int cell;
};

// Returns an expression encoding the possible combinations of candidates which
// can be selected.
std::vector<Token> BuildExpression(
    const Subdivision& subdivision,
    const std::multimap<int, const Candidate*>& candidates_per_cell) {
  // Start from the Subdivision root (cell = 0), and do a depth-first traversal
  // to emit an expression which encodes the following constraint:
  //   Given a cell, either select one of the candidates for that cell, or
  //   select all of the options for all child cells.
  //
  // For example, consider the very simple Subdivision with 3 cells:
  //    0
  //   / \
  //  1   2
  // Where:
  //  cell 0 has candidates a, b.
  //  cell 1 has candidates c, d, e.
  //  cell 2 has candidate f.
  //
  // The resulting expression will be:
  //   And(Or(a, b, And(Or(c, d, e), Or(f))))
  //
  // Note that the outer-most And() clause is typically redundant, but it
  // ensures that a well-formed (i.e. non-empty) expression is sent to the
  // solver in all cases.
  std::vector<Token> expression;
  const int kEndOfChildMarker = -1;
  std::vector<int> cell_to_visit;
  cell_to_visit.push_back(kEndOfChildMarker);
  subdivision.GetRoots(&cell_to_visit);
  expression.push_back(Token::And());
  while (!cell_to_visit.empty()) {
    int cell = cell_to_visit.back();
    cell_to_visit.pop_back();

    if (cell == kEndOfChildMarker) {
      expression.push_back(Token::End());
      expression.push_back(Token::End());
      continue;
    }

    expression.push_back(Token::Or());
    auto candidate_range = candidates_per_cell.equal_range(cell);
    if (candidate_range.first != candidates_per_cell.end()) {
      for (auto candidate_iter = candidate_range.first;
           candidate_iter != candidate_range.second; ++candidate_iter) {
        expression.push_back(Token::Item(candidate_iter->second->item));
      }
    }
    cell_to_visit.push_back(kEndOfChildMarker);
    subdivision.GetChildren(cell, &cell_to_visit);
    if (cell_to_visit.back() == kEndOfChildMarker) {
      // There were no children.
      cell_to_visit.pop_back();
      expression.push_back(Token::End());
      continue;
    }
    expression.push_back(Token::And());
  }
  expression.push_back(Token::End());
  return expression;
}

}  // namespace

std::vector<Tile> SelectionTiler::Run(const PointSet& points) {
  LOG(INFO) << "Running tiler with points: " << points.positions.size();
  subdivision_->Init(points);
  candidate_generator_->Init(points);

  const int num_weights = tile_weight_model_->GetDimension();

  std::vector<int> cells;
  GetCellsInDepthRange(*subdivision_, min_subdivision_level_,
                       std::numeric_limits<int>::max(), &cells);

  // Generate all candidate tilings for all relevant pyramid cells.
  std::vector<std::vector<CandidateTileGenerator::CandidateTiles>>
      candidate_tiles_per_cell(cells.size());
  candidate_generator_->GenerateCandidateTiles(
      points, *subdivision_, cells, absl::MakeSpan(candidate_tiles_per_cell));

  // Use candidate_tiles_per_cell to generate the ItemSet and mapping between
  // cells, candidate tiles, and items.
  selection::ItemSet item_set(num_weights);
  std::vector<Candidate> all_candidates;
  {
    std::vector<float> dense_weights_for_tile(num_weights);
    std::vector<float> dense_weights(num_weights);
    std::vector<selection::ItemSet::Weight> sparse_weights;
    sparse_weights.reserve(num_weights);
    for (int i = 0; i < cells.size(); ++i) {
      int cell = cells[i];
      for (const auto& tile_set : candidate_tiles_per_cell[i]) {
        float total_cost =
            std::accumulate(tile_set.costs.begin(), tile_set.costs.end(), 0.0f);
        std::fill(dense_weights_for_tile.begin(), dense_weights_for_tile.end(),
                  0.0f);
        std::fill(dense_weights.begin(), dense_weights.end(), 0.0f);
        for (const auto& tile : tile_set.tiles) {
          tile_weight_model_->GetTileWeight(
              tile, absl::MakeSpan(dense_weights_for_tile));
          for (int w = 0; w < num_weights; ++w) {
            dense_weights[w] += dense_weights_for_tile[w];
          }
        }
        sparse_weights.clear();
        for (int w = 0; w < num_weights; ++w) {
          if (dense_weights[w] > 0.0f) {
            // Normalize weights.
            float normalized_weight = dense_weights[w] / max_weight_[w];
            sparse_weights.push_back({w, normalized_weight});
          }
        }
        int item = item_set.AppendItem(total_cost, sparse_weights);
        DCHECK_EQ(all_candidates.size(), item);
        all_candidates.push_back({&tile_set, item, cell});
      }
    }
  }

  std::multimap<int, const Candidate*> candidates_per_cell;
  for (const auto& candidate : all_candidates) {
    candidates_per_cell.insert({candidate.cell, &candidate});
  }
  std::vector<Token> expression =
      BuildExpression(*subdivision_, candidates_per_cell);

  // Capacity is 1 since weights have been normalized above.
  std::vector<double> capacity(num_weights, 1.0);

  selection::SelectionProblem problem;
  problem.items = &item_set;
  problem.capacity = capacity;
  problem.expression = expression;

  std::vector<int> solution_items;

  selection_solver_->Init(problem);
  std::vector<double> multipliers(num_weights, 0.0);
  bool success =
      selection_solver_->Solve(absl::MakeSpan(multipliers), &solution_items);
  if (!success) {
    base::SeuratError(
        "Failed to generate geometry within specified constraints.");
  }

  float total_cost = 0.0f;
  std::vector<Tile> tiles;
  for (int item : solution_items) {
    const auto& tiles_for_item = all_candidates[item].tiles->tiles;
    tiles.insert(tiles.end(), tiles_for_item.begin(), tiles_for_item.end());
    total_cost +=
        std::accumulate(all_candidates[item].tiles->costs.begin(),
                        all_candidates[item].tiles->costs.end(), 0.0f);
  }
  LOG(INFO) << "Solver completed with total cost: " << total_cost;
  return tiles;
}

}  // namespace tiler
}  // namespace seurat
